/** @file  NorFlashDxe.c

  Copyright (c) 2011 - 2014, ARM Ltd. All rights reserved.<BR>
  Copyright (C) 2022 - 2023, Phytium Technology Co., Ltd. All rights reserved.<BR>

  This program and the accompanying materials
  are licensed and made available under the terms and conditions of the BSD License
  which accompanies this distribution.  The full text of the license may be found at
  http://opensource.org/licenses/bsd-license.php

  THE PROGRAM IS DISTRIBUTED UNDER THE BSD LICENSE ON AN "AS IS" BASIS,
  WITHOUT WARRANTIES OR REPRESENTATIONS OF ANY KIND, EITHER EXPRESS OR IMPLIED.

**/

#include <Library/UefiLib.h>
#include <Library/BaseMemoryLib.h>
#include <Library/MemoryAllocationLib.h>
#include <Library/UefiBootServicesTableLib.h>
#include <Library/PcdLib.h>
#include <Library/DxeServicesTableLib.h>
#include <Library/PhytiumSpiNorFlashLib.h>


#include "NorFlashDxe.h"


STATIC EFI_EVENT mNorFlashVirtualAddrChangeEvent;


//
// Global variable declarations
//
NOR_FLASH_INSTANCE  **mNorFlashInstances;
UINTN               *mNorFlashDeviceCount;


NOR_FLASH_INSTANCE  mNorFlashInstanceTemplate = {
  NOR_FLASH_SIGNATURE, // Signature
  NULL, // Handle ... NEED TO BE FILLED

  FALSE, // Initialized
  NULL, // Initialize

  0, // DeviceBaseAddress ... NEED TO BE FILLED
  0, // RegionBaseAddress ... NEED TO BE FILLED
  0, // Size ... NEED TO BE FILLED
  0, // StartLba

  {
    EFI_BLOCK_IO_PROTOCOL_REVISION2, // Revision
    NULL, // Media ... NEED TO BE FILLED
    NorFlashBlockIoReset, // Reset;
    NorFlashBlockIoReadBlocks,          // ReadBlocks
    NorFlashBlockIoWriteBlocks,         // WriteBlocks
    NorFlashBlockIoFlushBlocks          // FlushBlocks
  }, // BlockIoProtocol

  {
    0, // MediaId ... NEED TO BE FILLED
    FALSE, // RemovableMedia
    TRUE, // MediaPresent
    FALSE, // LogicalPartition
    FALSE, // ReadOnly
    FALSE, // WriteCaching;
    0, // BlockSize ... NEED TO BE FILLED
    4, //  IoAlign
    0, // LastBlock ... NEED TO BE FILLED
    0, // LowestAlignedLba
    1, // LogicalBlocksPerPhysicalBlock
  }, //Media;

  {
    EFI_DISK_IO_PROTOCOL_REVISION, // Revision
    NorFlashDiskIoReadDisk,        // ReadDisk
    NorFlashDiskIoWriteDisk        // WriteDisk
  },

  FALSE, // SupportFvb ... NEED TO BE FILLED
  {
    FvbGetAttributes, // GetAttributes
    FvbSetAttributes, // SetAttributes
    FvbGetPhysicalAddress,  // GetPhysicalAddress
    FvbGetBlockSize,  // GetBlockSize
    FvbRead,  // Read
    FvbWrite, // Write
    FvbEraseBlocks, // EraseBlocks
    NULL, //ParentHandle
  }, //  FvbProtoccol;

  NULL, // ShadowBuffer
  {
    {
      {
        HARDWARE_DEVICE_PATH,
        HW_VENDOR_DP,
        { (UINT8)sizeof(VENDOR_DEVICE_PATH), (UINT8)((sizeof(VENDOR_DEVICE_PATH)) >> 8) }
      },
      { 0x0, 0x0, 0x0, { 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0 } }, // GUID ... NEED TO BE FILLED
    },
    {
      END_DEVICE_PATH_TYPE,
      END_ENTIRE_DEVICE_PATH_SUBTYPE,
      { sizeof (EFI_DEVICE_PATH_PROTOCOL), 0 }
    }
    } // DevicePath
};


EFI_STATUS
NorFlashCreateInstance (
  IN UINTN                  NorFlashDeviceBase,
  IN UINTN                  NorFlashRegionBase,
  IN UINTN                  NorFlashSize,
  IN UINT32                 MediaId,
  IN UINT32                 BlockSize,
  IN BOOLEAN                SupportFvb,
  IN CONST GUID            *NorFlashGuid,
  OUT NOR_FLASH_INSTANCE**  NorFlashInstance
  )
{
  EFI_STATUS Status;
  NOR_FLASH_INSTANCE* Instance;

  ASSERT(NorFlashInstance != NULL);

  Instance = AllocateRuntimeCopyPool (sizeof(NOR_FLASH_INSTANCE),&mNorFlashInstanceTemplate);
  if (Instance == NULL) {
    return EFI_OUT_OF_RESOURCES;
  }

  Instance->DeviceBaseAddress = NorFlashDeviceBase;
  Instance->RegionBaseAddress = NorFlashRegionBase;
  Instance->Size = NorFlashSize;

  Instance->BlockIoProtocol.Media = &Instance->Media;
  Instance->Media.MediaId = MediaId;
  Instance->Media.BlockSize = BlockSize;
  Instance->Media.LastBlock = (NorFlashSize / BlockSize)-1;

  CopyGuid (&Instance->DevicePath.Vendor.Guid, NorFlashGuid);

  Instance->ShadowBuffer = AllocateRuntimePool (BlockSize);;
  if (Instance->ShadowBuffer == NULL) {
    return EFI_OUT_OF_RESOURCES;
  }

  if (SupportFvb) {
    Instance->SupportFvb = TRUE;
    Instance->Initialize = NorFlashFvbInitialize;

    Status = gBS->InstallMultipleProtocolInterfaces (
                  &Instance->Handle,
                  &gEfiDevicePathProtocolGuid, &Instance->DevicePath,
                  // &gEfiBlockIoProtocolGuid, &Instance->BlockIoProtocol,
                  &gEfiFirmwareVolumeBlockProtocolGuid, &Instance->FvbProtocol,
                  NULL
                  );
    if (EFI_ERROR(Status)) {
      FreePool (Instance);
      return Status;
    }
  } else {
    Instance->Initialized = TRUE;

    Status = gBS->InstallMultipleProtocolInterfaces (
                  &Instance->Handle,
                  &gEfiDevicePathProtocolGuid, &Instance->DevicePath,
                  // &gEfiBlockIoProtocolGuid, &Instance->BlockIoProtocol,
                  &gEfiDiskIoProtocolGuid, &Instance->DiskIoProtocol,
                  NULL
                  );
    if (EFI_ERROR(Status)) {
      FreePool (Instance);
      return Status;
    }
  }

  *NorFlashInstance = Instance;
  return Status;
}

STATIC
BOOLEAN
NorFlashBlockIsLocked (
  IN NOR_FLASH_INSTANCE     *Instance,
  IN UINTN                  BlockAddress
  )
{
  return FALSE;
}

STATIC
EFI_STATUS
NorFlashUnlockSingleBlock (
  IN NOR_FLASH_INSTANCE     *Instance,
  IN UINTN                  BlockAddress
  )
{
  DEBUG((EFI_D_BLKIO, "UnlockSingleBlock: BlockAddress=0x%08x\n", BlockAddress));

  return EFI_SUCCESS;
}

STATIC
EFI_STATUS
NorFlashUnlockSingleBlockIfNecessary (
  IN NOR_FLASH_INSTANCE     *Instance,
  IN UINTN                  BlockAddress
  )
{
  EFI_STATUS Status;

  Status = EFI_SUCCESS;

  if (NorFlashBlockIsLocked (Instance, BlockAddress) == TRUE) {
    Status = NorFlashUnlockSingleBlock (Instance, BlockAddress);
  }

  return Status;
}

/**
 * This function unlock and erase an entire NOR Flash block.
 **/
EFI_STATUS
NorFlashUnlockAndEraseSingleBlock (
  IN NOR_FLASH_INSTANCE     *Instance,
  IN UINTN                  BlockAddress
  )
{
  EFI_STATUS      Status;
  UINTN           Index;
  EFI_TPL         OriginalTPL;

  if (!EfiAtRuntime ()) {
    // Raise TPL to TPL_HIGH to stop anyone from interrupting us.
    OriginalTPL = gBS->RaiseTPL (TPL_HIGH_LEVEL);
  } else {
    // This initialization is only to prevent the compiler to complain about the
    // use of uninitialized variables
    OriginalTPL = TPL_HIGH_LEVEL;
  }

  Index = 0;
  // The block erase might fail a first time (SW bug ?). Retry it ...
  do {
    // Unlock the block if we have to
    Status = NorFlashUnlockSingleBlockIfNecessary (Instance, BlockAddress);
    if (EFI_ERROR (Status)) {
      break;
    }
    Status = NorFlashPlatformEraseSingleBlock(BlockAddress);
    Index++;
  } while ((Index < NOR_FLASH_ERASE_RETRY) && (Status == EFI_WRITE_PROTECTED));

  if (Index == NOR_FLASH_ERASE_RETRY) {
    DEBUG((EFI_D_ERROR,"EraseSingleBlock(BlockAddress=0x%08x: Block Locked Error (try to erase %d times)\n", BlockAddress,Index));
  }

  if (!EfiAtRuntime ()) {
    // Interruptions can resume.
    gBS->RestoreTPL (OriginalTPL);
  }

  return Status;
}

EFI_STATUS
NorFlashRead (
  IN NOR_FLASH_INSTANCE   *Instance,
  IN EFI_LBA               Lba,
  IN UINTN                 Offset,
  IN UINTN                 BufferSizeInBytes,
  OUT VOID                *Buffer
  )
{
  UINTN Address = GET_NOR_BLOCK_ADDRESS(Instance->RegionBaseAddress, Lba, Instance->Media.BlockSize)  + Offset;

  // if not initialized, initialize instance first
  if (!Instance->Initialized && Instance->Initialize) {
    Instance->Initialize(Instance);
  }

  if (BufferSizeInBytes == 0) {
    return EFI_SUCCESS;
  }

  // The buffer must be valid
  if (Buffer == NULL) {
    return EFI_INVALID_PARAMETER;
  }

  return NorFlashPlatformRead(Address, Buffer, BufferSizeInBytes);
}

/*
 * Write a full or portion of a block. It must not span block boundaries; that is,
 * Offset + *NumBytes <= Instance->Media.BlockSize.
 */
EFI_STATUS
NorFlashWriteSingleBlock (
  IN NOR_FLASH_INSTANCE   *Instance,
  IN EFI_LBA               Lba,
  IN UINTN                 Offset,
  IN UINTN                 BufferSizeInBytes,
  IN UINT8                *Buffer
  )
{
  EFI_STATUS  Status;
  UINT32      Tmp;
  UINT32      TmpBuf;
  UINT32      WordToWrite;
  UINT32      Mask;
  BOOLEAN     DoErase;
  UINTN       BytesToWrite;
  UINTN       CurOffset;
  UINTN       WordAddr;
  UINTN       BlockSize;
  UINTN       BlockAddress;

  if (!Instance->Initialized && Instance->Initialize) {
    Instance->Initialize(Instance);
  }

  // Detect WriteDisabled state
  if (Instance->Media.ReadOnly == TRUE) {
    DEBUG ((EFI_D_ERROR, "NorFlashWriteSingleBlock: ERROR - Can not write: Device is in WriteDisabled state.\n"));
    // It is in WriteDisabled state, return an error right away
    return EFI_ACCESS_DENIED;
  }

  // Cache the block size to avoid de-referencing pointers all the time
  BlockSize = Instance->Media.BlockSize;

  // The write must not span block boundaries.
  // We need to check each variable individually because adding two large values together overflows.
  if ((Offset                       >= BlockSize) ||
      (BufferSizeInBytes            >  BlockSize) ||
      ((Offset + BufferSizeInBytes) >  BlockSize)) {
    DEBUG ((EFI_D_ERROR, "NorFlashWriteSingleBlock: ERROR - EFI_BAD_BUFFER_SIZE: (Offset=0x%x + NumBytes=0x%x) > BlockSize=0x%x\n", Offset, BufferSizeInBytes, BlockSize ));
    return EFI_BAD_BUFFER_SIZE;
  }

  // We must have some bytes to write
  if (BufferSizeInBytes == 0) {
    DEBUG ((EFI_D_ERROR, "NorFlashWriteSingleBlock: ERROR - EFI_BAD_BUFFER_SIZE: NumBytes == 0\n"));
    return EFI_BAD_BUFFER_SIZE;
  }

  // Pick 128bytes as a good start for word operations as opposed to erasing the
  // block and writing the data regardless if an erase is really needed.
  // It looks like most individual NV variable writes are smaller than 128bytes.
  if (BufferSizeInBytes <= 128) {
    // Check to see if we need to erase before programming the data into NOR.
    // If the destination bits are only changing from 1s to 0s we can just write.
    // After a block is erased all bits in the block is set to 1.
    // If any byte requires us to erase we just give up and rewrite all of it.
    DoErase      = FALSE;
    BytesToWrite = BufferSizeInBytes;
    CurOffset    = Offset;

    while (BytesToWrite > 0) {
      // Read full word from NOR, splice as required. A word is the smallest
      // unit we can write.
      Status = NorFlashRead (Instance, Lba, CurOffset & ~(0x3), sizeof(Tmp), &Tmp);
      if (EFI_ERROR (Status)) {
        return EFI_DEVICE_ERROR;
      }

      // Physical address of word in NOR to write.
      WordAddr = (CurOffset & ~(0x3)) + GET_NOR_BLOCK_ADDRESS (Instance->RegionBaseAddress,
                                                               Lba, BlockSize);
      // The word of data that is to be written.
      TmpBuf = *((UINT32*)(Buffer + (BufferSizeInBytes - BytesToWrite)));

      // First do word aligned chunks.
      if ((CurOffset & 0x3) == 0) {
        if (BytesToWrite >= 4) {
          // Is the destination still in 'erased' state?
          if (~Tmp != 0) {
            // Check to see if we are only changing bits to zero.
            if ((Tmp ^ TmpBuf) & TmpBuf) {
              DoErase = TRUE;
              break;
            }
          }
          // Write this word to NOR
          WordToWrite = TmpBuf;
          CurOffset += sizeof(TmpBuf);
          BytesToWrite -= sizeof(TmpBuf);
        } else {
          // BytesToWrite < 4. Do small writes and left-overs
          Mask = ~((~0) << (BytesToWrite * 8));
          // Mask out the bytes we want.
          TmpBuf &= Mask;
          // Is the destination still in 'erased' state?
          if ((Tmp & Mask) != Mask) {
            // Check to see if we are only changing bits to zero.
            if ((Tmp ^ TmpBuf) & TmpBuf) {
              DoErase = TRUE;
              break;
            }
          }
          // Merge old and new data. Write merged word to NOR
          WordToWrite = (Tmp & ~Mask) | TmpBuf;
          CurOffset += BytesToWrite;
          BytesToWrite = 0;
        }
      } else {
        // Do multiple words, but starting unaligned.
        if (BytesToWrite > (4 - (CurOffset & 0x3))) {
          Mask = ((~0) << ((CurOffset & 0x3) * 8));
          // Mask out the bytes we want.
          TmpBuf &= Mask;
          // Is the destination still in 'erased' state?
          if ((Tmp & Mask) != Mask) {
            // Check to see if we are only changing bits to zero.
            if ((Tmp ^ TmpBuf) & TmpBuf) {
              DoErase = TRUE;
              break;
            }
          }
          // Merge old and new data. Write merged word to NOR
          WordToWrite = (Tmp & ~Mask) | TmpBuf;
          BytesToWrite -= (4 - (CurOffset & 0x3));
          CurOffset += (4 - (CurOffset & 0x3));
        } else {
          // Unaligned and fits in one word.
          Mask = (~((~0) << (BytesToWrite * 8))) << ((CurOffset & 0x3) * 8);
          // Mask out the bytes we want.
          TmpBuf = (TmpBuf << ((CurOffset & 0x3) * 8)) & Mask;
          // Is the destination still in 'erased' state?
          if ((Tmp & Mask) != Mask) {
            // Check to see if we are only changing bits to zero.
            if ((Tmp ^ TmpBuf) & TmpBuf) {
              DoErase = TRUE;
              break;
            }
          }
          // Merge old and new data. Write merged word to NOR
          WordToWrite = (Tmp & ~Mask) | TmpBuf;
          CurOffset += BytesToWrite;
          BytesToWrite = 0;
        }
      }

      //
      // Write the word to NOR.
      //
      //[jwluo-0015-start]
      //Status = NorFlashPlatformWrite (WordAddr, &WordToWrite, 4);
        Status = NorFlashPlatformWriteAtRunTime (Instance->RegionBaseAddress, WordAddr, &WordToWrite, 4);
      //[jwluo-0015-end]
    if (EFI_ERROR (Status)) {
        return EFI_DEVICE_ERROR;
      }
    }
    // Exit if we got here and could write all the data. Otherwise do the
    // Erase-Write cycle.
    if (!DoErase) {
      return EFI_SUCCESS;
    }
  }

  // Check we did get some memory. Buffer is BlockSize.
  if (Instance->ShadowBuffer == NULL) {
    DEBUG ((EFI_D_ERROR, "NorFlashWriteSingleBlock: ERROR - Buffer not ready\n"));
    return EFI_DEVICE_ERROR;
  }

  DEBUG ((EFI_D_BLKIO, "NorFlashWriteSingleBlock: (Lba=%ld, Offset=0x%x, NumBytes=0x%x, Buffer @ 0x%08x)\n", Lba, Offset, BufferSizeInBytes, Buffer));
  BlockAddress = GET_NOR_BLOCK_ADDRESS(Instance->RegionBaseAddress, Lba, Instance->Media.BlockSize);
  // Read NOR Flash data into shadow buffer
  Status = NorFlashPlatformRead (BlockAddress, Instance->ShadowBuffer, BlockSize);
  if (EFI_ERROR (Status)) {
    // Return one of the pre-approved error statuses
    return EFI_DEVICE_ERROR;
  }

  // Put the data at the appropriate location inside the buffer area
  CopyMem ((VOID*)((UINTN)Instance->ShadowBuffer + Offset), Buffer, BufferSizeInBytes);

  // Write the modified buffer back to the NorFlash
  //[-START-KLK-Kevin-0013-20211210-MODIFY]// 
  //Status = NorFlashPlatformEraseSingleBlock (BlockAddress);
  Status = NorFlashPlatformEraseSingleBlock (BlockAddress - Instance->RegionBaseAddress);
  //[-END-KLK-Kevin-0013-20211210-MODIFY]// 
  if (EFI_ERROR (Status)) {
    // Return one of the pre-approved error statuses
    return EFI_DEVICE_ERROR;
  }

  // Write the modified buffer back to the NorFlash
  //[jwluo-0015-start]
  //Status = NorFlashPlatformWrite (BlockAddress, Instance->ShadowBuffer, BlockSize);
    Status = NorFlashPlatformWriteAtRunTime (Instance->RegionBaseAddress, BlockAddress, Instance->ShadowBuffer, BlockSize);
  //[jwluo-0015-end]
  if (EFI_ERROR (Status)) {
    // Return one of the pre-approved error statuses
    return EFI_DEVICE_ERROR;
  }

  return EFI_SUCCESS;
}

EFI_STATUS
NorFlashWrite (
  IN NOR_FLASH_INSTANCE    *Instance,
  IN EFI_LBA                Lba,
  IN UINTN                  Offset,
  IN UINTN                  BufferSizeInBytes,
  IN VOID                  *Buffer
  )
{
  EFI_STATUS      Status = EFI_SUCCESS;
  UINT32          BlockSize;
  UINT32          BlockOffset;
  UINTN           RemainingBytes;
  UINTN           WriteSize;

  // if not initialized, initialize instance first
  if (!Instance->Initialized && Instance->Initialize) {
    Instance->Initialize(Instance);
  }

  if (Instance->Media.ReadOnly == TRUE) {
    return EFI_WRITE_PROTECTED;
  }

  if (BufferSizeInBytes == 0) {
    return EFI_SUCCESS;
  }

  if (Buffer == NULL) {
    return EFI_INVALID_PARAMETER;
  }

  BlockSize      = Instance->Media.BlockSize;
  BlockOffset    = Offset;
  RemainingBytes = BufferSizeInBytes;

  // The write must not span block boundaries.
  // We need to check each variable individually because adding two large values together overflows.
  if (Offset >= BlockSize) {
    DEBUG ((EFI_D_ERROR, "NorFlashWrite: ERROR - EFI_BAD_BUFFER_SIZE: Offset=0x%x > BlockSize=0x%x\n", Offset, BlockSize));
    return EFI_BAD_BUFFER_SIZE;
  }

  // We must have some bytes to read
  DEBUG((EFI_D_BLKIO, "NorFlashWrite: Lba=0x%x Offset=0x%x BufferSizeInBytes=0x%x Buffer:0x%p\n",
        Lba, Offset, BufferSizeInBytes, Buffer));

  // Write either all the remaining bytes, or the number of bytes that bring
  // us up to a block boundary, whichever is less.
  // (DiskOffset | (BlockSize - 1)) + 1) rounds DiskOffset up to the next
  // block boundary (even if it is already on one).
  WriteSize = MIN (RemainingBytes, BlockSize - BlockOffset);

  do {
    Status = NorFlashWriteSingleBlock (Instance, Lba, BlockOffset, WriteSize, Buffer);
    if (EFI_ERROR (Status)) {
      return Status;
    }

    // Now continue writing either all the remaining bytes or single blocks.
    RemainingBytes -= WriteSize;
    Buffer = (UINT8 *) Buffer + WriteSize;
    Lba++;
    BlockOffset = 0;
    WriteSize = MIN (RemainingBytes, BlockSize);
  } while (RemainingBytes);

  return Status;
}

EFI_STATUS
NorFlashReset (
  IN NOR_FLASH_INSTANCE   *Instance
  )
{
  return EFI_SUCCESS;
}


/*
  Although DiskIoDxe will automatically install the DiskIO protocol whenever
  we install the BlockIO protocol, its implementation is sub-optimal as it reads
  and writes entire blocks using the BlockIO protocol. In fact we can access
  NOR flash with a finer granularity than that, so we can improve performance
  by directly producing the DiskIO protocol.
*/

/**
  Read BufferSize bytes from Offset into Buffer.

  @param  This                  Protocol instance pointer.
  @param  MediaId               Id of the media, changes every time the media is replaced.
  @param  Offset                The starting byte offset to read from
  @param  BufferSize            Size of Buffer
  @param  Buffer                Buffer containing read data

  @retval EFI_SUCCESS           The data was read correctly from the device.
  @retval EFI_DEVICE_ERROR      The device reported an error while performing the read.
  @retval EFI_NO_MEDIA          There is no media in the device.
  @retval EFI_MEDIA_CHNAGED     The MediaId does not matched the current device.
  @retval EFI_INVALID_PARAMETER The read request contains device addresses that are not
                                valid for the device.
**/
EFI_STATUS
EFIAPI
NorFlashDiskIoReadDisk (
  IN EFI_DISK_IO_PROTOCOL      *This,
  IN UINT32                     MediaId,
  IN UINT64                     DiskOffset,
  IN UINTN                      BufferSize,
  OUT VOID                     *Buffer
  )
{
  NOR_FLASH_INSTANCE *Instance;
  UINT32              BlockSize;
  UINT32              BlockOffset;
  EFI_LBA             Lba;

  Instance = INSTANCE_FROM_DISKIO_THIS(This);

  if (MediaId != Instance->Media.MediaId) {
    return EFI_MEDIA_CHANGED;
  }

  BlockSize = Instance->Media.BlockSize;
  Lba = (EFI_LBA) DivU64x32Remainder (DiskOffset, BlockSize, &BlockOffset);

  return NorFlashRead (Instance, Lba, BlockOffset, BufferSize, Buffer);
}

/**
  Writes a specified number of bytes to a device.

  @param  This       Indicates a pointer to the calling context.
  @param  MediaId    ID of the medium to be written.
  @param  Offset     The starting byte offset on the logical block I/O device to write.
  @param  BufferSize The size in bytes of Buffer. The number of bytes to write to the device.
  @param  Buffer     A pointer to the buffer containing the data to be written.

  @retval EFI_SUCCESS           The data was written correctly to the device.
  @retval EFI_WRITE_PROTECTED   The device can not be written to.
  @retval EFI_DEVICE_ERROR      The device reported an error while performing the write.
  @retval EFI_NO_MEDIA          There is no media in the device.
  @retval EFI_MEDIA_CHNAGED     The MediaId does not matched the current device.
  @retval EFI_INVALID_PARAMETER The write request contains device addresses that are not
                                valid for the device.
**/
EFI_STATUS
EFIAPI
NorFlashDiskIoWriteDisk (
  IN EFI_DISK_IO_PROTOCOL      *This,
  IN UINT32                     MediaId,
  IN UINT64                     DiskOffset,
  IN UINTN                      BufferSize,
  IN VOID                      *Buffer
  )
{
  NOR_FLASH_INSTANCE *Instance;
  UINT32              BlockSize;
  UINT32              BlockOffset;
  EFI_LBA             Lba;

  Instance = INSTANCE_FROM_DISKIO_THIS(This);

  if (MediaId != Instance->Media.MediaId) {
    return EFI_MEDIA_CHANGED;
  }

  BlockSize = Instance->Media.BlockSize;
  Lba = (EFI_LBA) DivU64x32Remainder (DiskOffset, BlockSize, &BlockOffset);

  return NorFlashWrite (Instance, Lba, BlockOffset, BufferSize, Buffer);
}


/**
  Fixup internal data so that EFI can be call in virtual mode.
  Call the passed in Child Notify event and convert any pointers in
  lib to virtual mode.

  @param[in]    Event   The Event that is being processed
  @param[in]    Context Event Context
**/
VOID
EFIAPI
NorFlashVirtualNotifyEvent (
  IN EFI_EVENT        Event,
  IN VOID             *Context
  )
{
  UINTN Index;

  for (Index = 0; Index < *mNorFlashDeviceCount; Index++) {
    if (!mNorFlashInstances[Index]->DeviceBaseAddress) {
      /* FIXME: EfiConvertPointer doesn't deal with zero address */
      mNorFlashInstances[Index]->DeviceBaseAddress |= 0x1;
    }
    if (!mNorFlashInstances[Index]->RegionBaseAddress) {
      /* FIXME: EfiConvertPointer doesn't deal with zero address */
      mNorFlashInstances[Index]->RegionBaseAddress |= 0x1;
    }

    EfiConvertPointer (0x0, (VOID**)&mNorFlashInstances[Index]->DeviceBaseAddress);
    EfiConvertPointer (0x0, (VOID**)&mNorFlashInstances[Index]->RegionBaseAddress);

    mNorFlashInstances[Index]->DeviceBaseAddress &= ~(EFI_PHYSICAL_ADDRESS)0x1;
    mNorFlashInstances[Index]->RegionBaseAddress &= ~(EFI_PHYSICAL_ADDRESS)0x1;

    // Convert BlockIo protocol
    EfiConvertPointer (0x0, (VOID**)&mNorFlashInstances[Index]->BlockIoProtocol.FlushBlocks);
    EfiConvertPointer (0x0, (VOID**)&mNorFlashInstances[Index]->BlockIoProtocol.ReadBlocks);
    EfiConvertPointer (0x0, (VOID**)&mNorFlashInstances[Index]->BlockIoProtocol.Reset);
    EfiConvertPointer (0x0, (VOID**)&mNorFlashInstances[Index]->BlockIoProtocol.WriteBlocks);

    // Convert Fvb
    if (mNorFlashInstances[Index]->SupportFvb) {
      EfiConvertPointer (0x0, (VOID**)&mNorFlashInstances[Index]->FvbProtocol.EraseBlocks);
      EfiConvertPointer (0x0, (VOID**)&mNorFlashInstances[Index]->FvbProtocol.GetAttributes);
      EfiConvertPointer (0x0, (VOID**)&mNorFlashInstances[Index]->FvbProtocol.GetBlockSize);
      EfiConvertPointer (0x0, (VOID**)&mNorFlashInstances[Index]->FvbProtocol.GetPhysicalAddress);
      EfiConvertPointer (0x0, (VOID**)&mNorFlashInstances[Index]->FvbProtocol.Read);
      EfiConvertPointer (0x0, (VOID**)&mNorFlashInstances[Index]->FvbProtocol.SetAttributes);
      EfiConvertPointer (0x0, (VOID**)&mNorFlashInstances[Index]->FvbProtocol.Write);
    } else {
      EfiConvertPointer (0x0, (VOID**)&mNorFlashInstances[Index]->DiskIoProtocol.ReadDisk);
      EfiConvertPointer (0x0, (VOID**)&mNorFlashInstances[Index]->DiskIoProtocol.WriteDisk);
    }

    if (mNorFlashInstances[Index]->ShadowBuffer != NULL) {
      EfiConvertPointer (0x0, (VOID**)&mNorFlashInstances[Index]->ShadowBuffer);
    }
  }

  EfiConvertPointer (0x0, (VOID **) &mNorFlashInstances);
  EfiConvertPointer (0x0, (VOID **) &mNorFlashDeviceCount);

  return;
}

extern const  UINT64  _gPcd_FixedAtBuild_PcdSpiControllerBase;
extern const  UINT64  _gPcd_FixedAtBuild_PcdSpiControllerSize;
#define SPI_CONTROLLER_BASE     _gPcd_FixedAtBuild_PcdSpiControllerBase
#define SPI_CONTROLLER_SIZE     _gPcd_FixedAtBuild_PcdSpiControllerSize

//[[jdzhang-0005]Add Begin
VOID
EFIAPI
SetMemorySpaceAttributes (
  IN EFI_EVENT        Event,
  IN VOID             *Context
  ){
  EFI_STATUS Status;
  Status = gDS->SetMemorySpaceAttributes (
      (SPI_CONTROLLER_BASE >> EFI_PAGE_SHIFT) << EFI_PAGE_SHIFT, EFI_PAGES_TO_SIZE(EFI_SIZE_TO_PAGES(SPI_CONTROLLER_SIZE)),
                  EFI_MEMORY_UC | EFI_MEMORY_RUNTIME
                  );
  DEBUG((EFI_D_ERROR, "[%a]:[%dL]%r \n", __FUNCTION__, __LINE__, Status));
  ASSERT_EFI_ERROR (Status);
}
//[[jdzhang-0005]Add End

EFI_STATUS
EFIAPI
NorFlashInitialise (
  IN EFI_HANDLE           ImageHandle,
  IN EFI_SYSTEM_TABLE    *SystemTable
  )
{
  EFI_STATUS              Status;
  UINT32                  Index;
  NOR_FLASH_DESCRIPTION_PHYTIUM*  NorFlashDevices;
  UINT32                  NorFlashDeviceCount;
  BOOLEAN                 ContainVariableStorage;
  //[[jdzhang-0005]Add Begin
  EFI_EVENT               Event;
  VOID                    *Registration;
  //[[jdzhang-0005]Add End

  Status = NorFlashPlatformInitialization ();
  if (EFI_ERROR(Status)) {
    DEBUG((EFI_D_ERROR,"NorFlashInitialise: Fail to initialize Nor Flash devices\n"));
    return Status;
  }
#if 1 // Moved to NorFlashInitialise
  //
  // Declare the SPI Controller Space as EFI_MEMORY_RUNTIME
  //
  Status = gDS->AddMemorySpace (
                  EfiGcdMemoryTypeMemoryMappedIo,
                  (SPI_CONTROLLER_BASE >> EFI_PAGE_SHIFT) << EFI_PAGE_SHIFT, EFI_PAGES_TO_SIZE(EFI_SIZE_TO_PAGES(SPI_CONTROLLER_SIZE)),
                  EFI_MEMORY_UC | EFI_MEMORY_RUNTIME
                  );
  DEBUG((EFI_D_ERROR, "[%a]:[%dL]%r \n", __FUNCTION__, __LINE__, Status));
  ASSERT_EFI_ERROR (Status);
  
  //[jdzhang-0005]Modify Begin
  //SetMemorySpaceAttributes fun will use gEfiCpuArchProtocolGuid, so do it in the protocol notify function.
  
  //Status = gDS->SetMemorySpaceAttributes (
  //  (SPI_CONTROLLER_BASE >> EFI_PAGE_SHIFT) << EFI_PAGE_SHIFT, EFI_PAGES_TO_SIZE(EFI_SIZE_TO_PAGES(SPI_CONTROLLER_SIZE)),
  //                EFI_MEMORY_UC | EFI_MEMORY_RUNTIME
  //                );
  //DEBUG((EFI_D_ERROR, "[%a]:[%dL]%r \n", __FUNCTION__, __LINE__, Status));
  //ASSERT_EFI_ERROR (Status);
  
   Status = gBS->CreateEvent (
                  EVT_NOTIFY_SIGNAL,
                  TPL_CALLBACK,
                  SetMemorySpaceAttributes,
                  NULL,
                  &Event
                  );
  ASSERT_EFI_ERROR (Status);

  //
  // Register for protocol notifications on this event
  //

  Status = gBS->RegisterProtocolNotify (
                  &gEfiCpuArchProtocolGuid,
                  Event,
                  &Registration
                  );

  ASSERT_EFI_ERROR (Status);
  //[jdzhang-0005]Modify End
  

#endif

  Status = NorFlashPlatformGetDevices (&NorFlashDevices, &NorFlashDeviceCount);
  if (EFI_ERROR(Status)) {
    DEBUG((EFI_D_ERROR,"NorFlashInitialise: Fail to get Nor Flash devices\n"));
    return Status;
  }

  mNorFlashInstances    = AllocateRuntimePool (sizeof(NOR_FLASH_INSTANCE*) * NorFlashDeviceCount);
  mNorFlashDeviceCount  = AllocateRuntimePool (sizeof(UINTN));
  *mNorFlashDeviceCount = NorFlashDeviceCount;

  for (Index = 0; Index < NorFlashDeviceCount; Index++) {
#if 0
    // Note: all the NOR Flash region needs to be reserved into the UEFI Runtime memory;
    //       even if we only use the small block region at the top of the NOR Flash.
    //       The reason is when the NOR Flash memory is set into program mode, the command
    //       is written as the base of the flash region (ie: Instance->DeviceBaseAddress)
    // Todo: SPI control block should be remapped, otherwise ...
    Status = gDS->AddMemorySpace (
	                  EfiGcdMemoryTypeMemoryMappedIo,
	                  NorFlashDevices[Index].RegionBaseAddress, NorFlashDevices[Index].Size,
	                  EFI_MEMORY_UC | EFI_MEMORY_RUNTIME
	                  );
    ASSERT_EFI_ERROR (Status);

    Status = gDS->SetMemorySpaceAttributes (
		    NorFlashDevices[Index].RegionBaseAddress, NorFlashDevices[Index].Size,
	                  EFI_MEMORY_UC | EFI_MEMORY_RUNTIME
	                  );
    ASSERT_EFI_ERROR (Status);
#endif
    // Check if this NOR Flash device contain the variable storage region
    ContainVariableStorage =
               (NorFlashDevices[Index].RegionBaseAddress <= PcdGet64 (PcdFlashNvStorageVariableBase64)) &&
               (PcdGet64 (PcdFlashNvStorageVariableBase64) + PcdGet32 (PcdFlashNvStorageVariableSize) <= NorFlashDevices[Index].RegionBaseAddress + NorFlashDevices[Index].Size);

    Status = NorFlashCreateInstance (
               NorFlashDevices[Index].DeviceBaseAddress,
               NorFlashDevices[Index].RegionBaseAddress,
               NorFlashDevices[Index].Size,
               Index,
               NorFlashDevices[Index].BlockSize,
               ContainVariableStorage,
               &NorFlashDevices[Index].Guid,
               &mNorFlashInstances[Index]
               );
    if (EFI_ERROR(Status)) {
      DEBUG((EFI_D_ERROR,"NorFlashInitialise: Fail to create instance for NorFlash[%d]\n",Index));
    }
  }

  //
  // Register for the virtual address change event
  //
  Status = gBS->CreateEventEx (
                  EVT_NOTIFY_SIGNAL,
                  TPL_NOTIFY,
                  NorFlashVirtualNotifyEvent,
                  NULL,
                  &gEfiEventVirtualAddressChangeGuid,
                  &mNorFlashVirtualAddrChangeEvent
                  );
  ASSERT_EFI_ERROR (Status);

  return Status;
}
